package com.migration.model.migration.scripts;

import com.migration.model.migration.AbstractMultiThreadMigrationScript;
import com.migration.model.migration.CustomExecutableMigrationScript;
import com.migration.model.newdb.RiskChainCause;
import com.migration.model.olddb.*;
import com.migration.service.RiskChainCauseService;
import com.migration.service.RiskChainDetailService;
import jakarta.persistence.EntityManager;
import jakarta.persistence.EntityManagerFactory;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 事故链-原因表(risk_chain_cause)迁移脚本
 */
@Slf4j
@Component
public class A024_RiskChainCauseMigrationScript extends AbstractMultiThreadMigrationScript implements CustomExecutableMigrationScript {

    @Autowired
    @Qualifier("oldEntityManagerFactory")
    private EntityManagerFactory oldEntityManagerFactory;

    @Autowired
    @Qualifier("newEntityManagerFactory")
    private EntityManagerFactory newEntityManagerFactory;

    @Autowired
    private RiskChainCauseService riskChainCauseService;

    @Autowired
    private RiskChainDetailService riskChainDetailService;

    /**
     * 构造函数，初始化迁移配置
     */
    public A024_RiskChainCauseMigrationScript() {
        super(
                "risk_chain_cause",
                "事故链-原因数据迁移",
                "ra_deviate_reason",
                "risk_chain_cause",
                24
        );
    }

    @Override
    protected String generateMigrationSql() {
        return "";
    }

    @Override
    public int executeCustomMigration(DataSource dataSource) {
        // 清空目标表
        JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
        jdbcTemplate.execute(getTruncateSql());

        // 执行自定义迁移
        return executeWithEntityProcessingMultiThread(oldEntityManagerFactory, newEntityManagerFactory);
    }

    /**
     * 构建reasonId到RaDeviateReason的映射
     *
     * @param oldEntityManager 旧数据库实体管理器
     * @param reasonIds        原因ID集合
     * @return 原因ID到原因实体的映射
     */
    private Map<Long, RaDeviateReason> buildReasonMap(EntityManager oldEntityManager, Set<Long> reasonIds) {
        Map<Long, RaDeviateReason> reasonMap = new HashMap<>();
        if (!reasonIds.isEmpty()) {
            jakarta.persistence.criteria.CriteriaBuilder cb = oldEntityManager.getCriteriaBuilder();
            jakarta.persistence.criteria.CriteriaQuery<RaDeviateReason> cq = cb.createQuery(RaDeviateReason.class);
            jakarta.persistence.criteria.Root<RaDeviateReason> root = cq.from(RaDeviateReason.class);
            cq.where(root.get("id").in(reasonIds));
            List<RaDeviateReason> reasonList = oldEntityManager.createQuery(cq).getResultList();

            for (RaDeviateReason reason : reasonList) {
                reasonMap.put(reason.getId(), reason);
            }
        }
        return reasonMap;
    }

    protected int executeWithEntityProcessingMultiThread(EntityManagerFactory oldEntityManagerFactory, EntityManagerFactory newEntityManagerFactory) {
        EntityManager oldEntityManager = oldEntityManagerFactory.createEntityManager();
        AtomicInteger rowsProcessed = createAtomicCounter();

        // 创建原因表行数统计器
        AtomicInteger causeRows = new AtomicInteger(0);

        try {
            // 1. 取得原始数据
            // 获取所有事故链详情数据
            List<RaFaultChainsDetail> detailList = riskChainDetailService.getAllFaultChainsDetails(oldEntityManager);

            // 获取所有原因数据
            List<RaDeviateReason> reasonList = riskChainCauseService.getAllReasons(oldEntityManager);

            // 获取所有reasonId列表
            List<Long> reasonIds = reasonList.stream()
                    .map(RaDeviateReason::getId)
                    .collect(Collectors.toList());

            // 构建原因ID与租户ID的映射关系
            Map<Long, Integer> reasonToTenantMap = riskChainCauseService.buildReasonToTenantMap(oldEntityManager, reasonIds);

            // 构建有后果的原因ID集合
            Set<Long> reasonWithConsequenceSet = riskChainCauseService.buildReasonWithConsequenceSet(detailList);

            /*********************************************** 多线程迁移数据 **************************************************/
            // 创建线程池和计数器
            ExecutorService executorService = createExecutorService();
            CountDownLatch latch = createCountDownLatch(1);

            // 迁移原因数据 (RiskChainCause)
            executorService.submit(() -> {
                try {
                    EntityManager newEntityManager = newEntityManagerFactory.createEntityManager();
                    try {
                        newEntityManager.getTransaction().begin();

                        // 按批次处理原因数据
                        int batchSize = Math.min(1000, reasonList.size());
                        for (int i = 0; i < reasonList.size(); i += batchSize) {
                            int toIndex = Math.min(i + batchSize, reasonList.size());
                            List<RaDeviateReason> batchReasons = reasonList.subList(i, toIndex);

                            List<RiskChainCause> causeBatch = riskChainCauseService.migrateCauseData(
                                oldEntityManager,
                                newEntityManager,
                                batchReasons,
                                reasonToTenantMap,
                                reasonWithConsequenceSet
                            );
                            int causeBatchSize = causeBatch.size();
                            causeRows.addAndGet(causeBatchSize);
                            rowsProcessed.addAndGet(causeBatchSize);
                        }

                        log.info("事故链-原因表(risk_chain_cause)迁移完成，共迁移 {} 行数据", causeRows.get());
                        newEntityManager.getTransaction().commit();
                    } catch (Exception e) {
                        if (newEntityManager.getTransaction().isActive()) {
                            newEntityManager.getTransaction().rollback();
                        }
                        throw e;
                    } finally {
                        newEntityManager.close();
                        latch.countDown();
                    }
                } catch (Exception e) {
                    log.error("事故链-原因表迁移异常", e);
                    throw new RuntimeException(e);
                }
            });

            // 等待所有任务完成
            waitForCompletion(latch, executorService);

            // 打印总结日志
            log.info("事故链-原因表迁移总结 ========================");
            log.info("事故链-原因表(risk_chain_cause)：{} 行", causeRows.get());
            log.info("总计迁移数据：{} 行", rowsProcessed.get());

        } catch (Exception e) {
            log.error("事故链-原因表迁移异常", e);
            throw e;
        } finally {
            oldEntityManager.close();
        }

        return rowsProcessed.get();
    }
}
